import cv2
import numpy as np
import os
import random

img_width_size = 320
img_height_size = 160
CONV_INPUT = "input_layer_input"
calib_batch_size = 32

def load_valid_data(data_path):
    img_list = []
    class_list = sorted(os.listdir(data_path))
    for class_name in class_list:
        print(class_name)
        image_list = os.listdir(data_path + class_name)
        random.shuffle(image_list)

        for idx in range(min(len(image_list), calib_batch_size // len(class_list) + 1)):
            img_path = data_path + class_name + "/" + image_list[idx]
            if img_path[-4:] != '.jpg':
                continue
            img = cv2.imread(img_path)  # BGR
            if not img is None:
                img = cv2.resize(img, (img_width_size, img_height_size), interpolation = cv2.INTER_LINEAR)
                img_list.append(img)
                      
    return img_list

dataset_valid_path = '/media/DataSets/DAC/train/'
validSet_images = load_valid_data(dataset_valid_path)
validSet_images = np.array(validSet_images)
def calib_input(iter):
    images = []
    for index in range(0, calib_batch_size):
        images.append(validSet_images[index])

    return {CONV_INPUT: images}

